airbyte_cdk/sources/declarative/partition_routers/grouping_partition_router.py (1)

12-24: Nitpick: Clarify mention of SinglePartitionRouter in the docstring.

In the docstring (line 20), the router is referred to as a SinglePartitionRouter, but the code uses PartitionRouter. Might be worth updating to maintain consistency and reduce confusion, wdyt?

airbyte_cdk/sources/declarative/declarative_component_schema.yaml (2)

3133-3141: Integration in SimpleRetriever: Inclusion of GroupingPartitionRouter
It looks like you added references to the new GroupingPartitionRouter inside the partition_router allowed types for SimpleRetriever. This should enable users to choose the grouping strategy for partitioning. Have you validated that the ordering and overall schema constraints are still correct here? wdyt?

---

3295-3302: Integration in AsyncRetriever: Inclusion of GroupingPartitionRouter
Similar to the SimpleRetriever, the AsyncRetriever now accepts GroupingPartitionRouter in its partition_router field. This is a great step toward consistency. Would you mind confirming that the asynchronous flow correctly processes grouped partitions as expected in practice? wdyt?

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py (1)

3052-3070: ⚠️ Potential issue

Fix parameter handling in GroupingPartitionRouter instantiation.

The pipeline failure indicates that parameters is not a valid argument for GroupingPartitionRouter. Looking at the implementation, we should remove this argument.

What do you think about this fix?

     return GroupingPartitionRouter(
         group_size=model.group_size,
         underlying_partition_router=underlying_router,
         deduplicate=model.deduplicate if model.deduplicate is not None else True,
-        parameters=model.parameters or {},
         config=config,
     )

airbyte_cdk/sources/declarative/models/declarative_component_schema.py (1)

2390-2411: 🛠️ Refactor suggestion

The GroupingPartitionRouter schema looks good, but should we update the type annotation?

The implementation looks solid with well-documented fields. However, there's a potential issue with the conint type annotation that might cause problems. Would you like to update it to use bracket notation instead?

-    group_size: conint(ge=1) = Field(
+    group_size: conint[ge=1] = Field(

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py (1)

3055-3057: Add type safety check for underlying_partition_router.

To prevent potential runtime errors, we should add a type check for the underlying_partition_router model.

What do you think about adding this type check?

+    if not isinstance(model.underlying_partition_router, BaseModel):
+        raise ValueError(
+            f"underlying_partition_router must be a BaseModel instance, got {type(model.underlying_partition_router)}"
+        )
+
     underlying_router = self._create_component_from_model(
         model=model.underlying_partition_router, config=config
     )

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

airbyte_cdk/sources/declarative/partition_routers/grouping_partition_router.py (2)

42-44: 🛠️ Refactor suggestion

Fix type safety in deduplication logic.

Based on past review comments, there are type safety issues in the deduplication logic. Would you consider this approach to improve type safety, wdyt?

-        batch = []
-        seen_keys = set()
+        from typing import List, Set, Any
+        batch: List[StreamSlice] = []
+        seen_keys: Set[Any] = set()

---

42-44: ⚠️ Potential issue

Fix type safety issues in deduplication logic.

The seen_keys variable needs type annotation to prevent type checker issues. Additionally, consider handling the case where partitions might be empty.

What do you think about these changes?

-        batch = []
-        seen_keys = set()
+        from typing import Set, Any
+        batch: list[StreamSlice] = []
+        seen_keys: Set[Any] = set()

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py (1)

3849-3903: LGTM! The test is well-structured and comprehensive.

The test effectively validates:

• Creation of GroupingPartitionRouter with correct group size
• Proper configuration of the underlying SubstreamPartitionRouter
• Parent stream configuration and field mappings

One suggestion: Would you consider adding a test case for deduplication behavior since GroupingPartitionRouter supports it? wdyt?

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py (2)

3053-3072: Add type safety checks for the underlying partition router.

The current implementation could be improved with better type safety. Would you consider adding type checks for model.underlying_partition_router before passing it to _create_component_from_model? This would help catch type-related issues earlier. For example:

 def create_grouping_partition_router(
     self, model: GroupingPartitionRouterModel, config: Config, **kwargs: Any
 ) -> GroupingPartitionRouter:
+    if not isinstance(model.underlying_partition_router, BaseModel):
+        raise ValueError(
+            f"underlying_partition_router must be a BaseModel instance, got {type(model.underlying_partition_router)}"
+        )
+
     underlying_router = self._create_component_from_model(
         model=model.underlying_partition_router, config=config
     )

wdyt?

---

3053-3072: Consider adding type hints for better type safety.

The implementation could benefit from explicit type hints. Would you consider these changes to improve type safety, wdyt?

    def create_grouping_partition_router(
        self, model: GroupingPartitionRouterModel, config: Config, **kwargs: Any
    ) -> GroupingPartitionRouter:
+       from typing import cast
        underlying_router = self._create_component_from_model(
            model=model.underlying_partition_router, config=config
        )
        if model.group_size < 1:
            raise ValueError(f"Group size must be greater than 0, got {model.group_size}")

        if not isinstance(underlying_router, PartitionRouter):
            raise ValueError(
                f"Underlying partition router must be a PartitionRouter subclass, got {type(underlying_router)}"
            )

        return GroupingPartitionRouter(
            group_size=model.group_size,
            underlying_partition_router=cast(PartitionRouter, underlying_router),
            deduplicate=model.deduplicate if model.deduplicate is not None else True,
            config=config,
        )

airbyte_cdk/sources/declarative/partition_routers/grouping_partition_router.py (4)

31-67: Consider adding error handling for empty partitions.

The current implementation assumes that partitions will always have values. Would you consider adding error handling for empty partitions? For example:

 def stream_slices(self) -> Iterable[StreamSlice]:
     batch = []
     seen_keys = set()

     for partition in self.underlying_partition_router.stream_slices():
+        if not partition.partition:
+            continue
+
         key = next(iter(partition.partition.values()), None)

Also, consider adding a docstring note about how empty partitions are handled, wdyt?

---

68-96: Consider handling edge cases in _create_grouped_slice.

The method assumes that all partitions in a batch have the same keys. Would you consider adding validation or handling for cases where partitions might have different keys? For example:

 def _create_grouped_slice(self, batch: list[StreamSlice]) -> StreamSlice:
+    if not batch:
+        raise ValueError("Cannot create a grouped slice from an empty batch")
+
     grouped_partition = {
-        key: [p.partition.get(key) for p in batch] for key in batch[0].partition.keys()
+        key: [p.partition.get(key) for p in batch] 
+        for key in set().union(*(p.partition.keys() for p in batch))
     }

wdyt?

---

68-97: Consider improving type hints in _create_grouped_slice method.

The method could benefit from more explicit type hints. Would you consider these changes, wdyt?

-    def _create_grouped_slice(self, batch: list[StreamSlice]) -> StreamSlice:
+    def _create_grouped_slice(self, batch: List[StreamSlice]) -> StreamSlice:
         """
         Creates a grouped StreamSlice from a batch of partitions, aggregating extra fields into a dictionary with list values.
 
         Args:
-            batch (list[StreamSlice]): A list of StreamSlice objects to group.
+            batch (List[StreamSlice]): A list of StreamSlice objects to group.
 
         Returns:
             StreamSlice: A single StreamSlice with combined partition and extra field values.
         """
         # Combine partition values into a single dict with lists
+        from typing import Dict
+        grouped_partition: Dict[str, list] = {
-        grouped_partition = {
             key: [p.partition.get(key) for p in batch] for key in batch[0].partition.keys()
         }
 
         # Aggregate extra fields into a dict with list values
-        extra_fields_dict = (
+        extra_fields_dict: Dict[str, list] = (

---

47-52: Consider adding error handling for empty partitions.

The code assumes that partition values exist. Would you consider adding validation to handle empty partitions gracefully, wdyt?

         # Extract the partition key (assuming single key-value pair, e.g., {"board_ids": value})
+        if not partition.partition:
+            continue
         key = next(iter(partition.partition.values()), None)
+        if key is None:
+            continue

airbyte_cdk/sources/declarative/declarative_component_schema.yaml (1)

3420-3457: Review the GroupingPartitionRouter schema definition:

• The new GroupingPartitionRouter component is clearly defined with its key properties: type, group_size, and underlying_partition_router. The renaming to underlying_partition_router nicely addresses the previous naming consistency comment.
• The description is comprehensive, explaining both the use case (grouping partitions for APIs supporting multiple keys in one request) and a potential pitfall with incremental syncs; this extra context is really helpful.
• The default value for the deduplicate flag is set to true, which seems sensible for ensuring unique partitions. Is this the intended behavior in most cases? Perhaps adding a brief note on why duplications are removed by default could further guide users.
• Would you consider providing a small example configuration snippet in the documentation? It might help users better understand how to integrate this new component into their stream settings.
  wdyt?

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py (1)

3053-3079: Consider adding type hints for the underlying_partition_router parameter.

The implementation looks solid with proper validation checks. However, we could improve type safety by adding type hints for the underlying_partition_router parameter in the model. What do you think about this change?

-        underlying_router = self._create_component_from_model(
-            model=model.underlying_partition_router, config=config
+        underlying_router: PartitionRouter = self._create_component_from_model(
+            model=model.underlying_partition_router, config=config
         )

Verify the behavior of deduplicate parameter.

The deduplicate parameter defaults to True if not specified. Could you confirm if this is the intended default behavior? Also, should we document this default value in the model's schema?

Consider adding validation for parent stream configurations.

The code checks for request options in parent stream configs but doesn't validate other potential issues. Should we add more comprehensive validation? For example:

1. Validate that parent streams don't have their own GroupingPartitionRouter to prevent nesting
2. Check for any other incompatible configurations
   wdyt?

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py (1)

3849-3999: Consider adding test cases for edge scenarios?

The current test coverage looks good but we could potentially add tests for:

• Deduplication behavior
• Different group sizes
• Multiple parent streams
• Empty parent stream results

WDYT about adding these test cases to make the test coverage more comprehensive?

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py (1)

3053-3082: Implementation looks good, but could we enhance the error messages?

The implementation properly validates group size and checks for unsupported configurations. I particularly like the comment explaining why request options aren't supported for GroupingPartitionRouter.

Just one suggestion - the error messages on lines 3071 and 3075 could be more descriptive to help users understand why request options don't work with grouped partitions, wdyt?

-                raise ValueError("Request options are not supported for GroupingPartitionRouter.")
+                raise ValueError("Request options are not supported for GroupingPartitionRouter because they are specific to individual partitions and can't be properly applied when partitions are grouped.")

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro